{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "LJ1byyLCzEea"
      },
      "source": [
        "Copyright 2022 Google LLC\n",
        "\n",
        "Licensed under the Apache License, Version 2.0 (the \"License\");\n",
        "you may not use this file except in compliance with the License.\n",
        "You may obtain a copy of the License at\n",
        "\n",
        "    https://www.apache.org/licenses/LICENSE-2.0\n",
        "\n",
        "Unless required by applicable law or agreed to in writing, software\n",
        "distributed under the License is distributed on an \"AS IS\" BASIS,\n",
        "WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
        "See the License for the specific language governing permissions and\n",
        "limitations under the License."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Bkq_0Rw8G9WY"
      },
      "source": [
        "# Colab notebook for demonstrating SCF calculations using the GAS22 functional\n",
        "\n",
        "SCF calculations are performed using PySCF package (v1.7.6). Functional derivatives are evaluated using automatic differentiation with JAX.\n",
        "\n",
        "Link to this notebook: https://colab.research.google.com/github/google-research/google-research/blob/master/symbolic_functionals/colab/run_GAS22.ipynb"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "executionInfo": {
          "elapsed": 14863,
          "status": "ok",
          "timestamp": 1644345778273,
          "user": {
            "displayName": "Li Li",
            "photoUrl": "https://lh3.googleusercontent.com/a-/AOh14GjUXuHpz1GDBTLkFU6msl96ZWWDxtMqu7UXHH6qLg=s64",
            "userId": "05582677413017445658"
          },
          "user_tz": 480
        },
        "id": "B_JlHcCODteJ",
        "outputId": "61502079-56b7-4455-c09c-6c310a8ec851"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Collecting pyscf==1.7.6\n",
            "  Downloading pyscf-1.7.6-cp37-cp37m-manylinux1_x86_64.whl (29.7 MB)\n",
            "\u001b[K     |████████████████████████████████| 29.7 MB 1.8 MB/s \n",
            "\u001b[?25hRequirement already satisfied: numpy!=1.16,!=1.17,\u003e1.8 in /usr/local/lib/python3.7/dist-packages (from pyscf==1.7.6) (1.19.5)\n",
            "Requirement already satisfied: scipy!=1.5.0,!=1.5.1 in /usr/local/lib/python3.7/dist-packages (from pyscf==1.7.6) (1.4.1)\n",
            "Requirement already satisfied: h5py\u003e2.2 in /usr/local/lib/python3.7/dist-packages (from pyscf==1.7.6) (3.1.0)\n",
            "Requirement already satisfied: cached-property in /usr/local/lib/python3.7/dist-packages (from h5py\u003e2.2-\u003epyscf==1.7.6) (1.5.2)\n",
            "Installing collected packages: pyscf\n",
            "Successfully installed pyscf-1.7.6\n"
          ]
        }
      ],
      "source": [
        "!pip install pyscf==1.7.6"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ZDQAYy9FEjoO"
      },
      "outputs": [],
      "source": [
        "import numpy as np\n",
        "import jax\n",
        "import jax.numpy as jnp\n",
        "from pyscf import dft\n",
        "from pyscf import gto\n",
        "\n",
        "jax.config.update('jax_enable_x64', True)\n",
        "\n",
        "HYBRID_COEFF = dft.libxc.hybrid_coeff('wb97mv')\n",
        "RSH_PARAMS = dft.libxc.rsh_coeff('wb97mv')"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "hEFvWs_YB1qE"
      },
      "source": [
        "# Helper functions to compute LDA densities and auxiliary quantities"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "HhL64ZkGqTgZ"
      },
      "outputs": [],
      "source": [
        "# _LDA_PW_PARAMETERS: dictionaries containing\n",
        "# parameters a, alpha, beta1, beta2, beta3, beta4 and fpp_0, which are used to\n",
        "# compute LDA correlation functional in the PW parametrization. The first six\n",
        "# parameters are used to evaluate auxiliary G functions in the PW paper\n",
        "# (10.1103/PhysRevB.45.13244). fpp_0 denotes the second derivative of\n",
        "# f_zeta function in the PW paper at zeta = 0.\n",
        "_LDA_PW_PARAMETERS = {\n",
        "    'eps_c_unpolarized': {\n",
        "        'a': 0.031091,\n",
        "        'alpha1': 0.21370,\n",
        "        'beta1': 7.5957,\n",
        "        'beta2': 3.5876,\n",
        "        'beta3': 1.6382,\n",
        "        'beta4': 0.49294},\n",
        "    'eps_c_polarized': {\n",
        "        'a': 0.015545,\n",
        "        'alpha1': 0.20548,\n",
        "        'beta1': 14.1189,\n",
        "        'beta2': 6.1977,\n",
        "        'beta3': 3.3662,\n",
        "        'beta4': 0.62517},\n",
        "    'alpha_c': {\n",
        "        'a': 0.016887,\n",
        "        'alpha1': 0.11125,\n",
        "        'beta1': 10.357,\n",
        "        'beta2': 3.6231,\n",
        "        'beta3': 0.88026,\n",
        "        'beta4': 0.49671},\n",
        "    'fpp_0': 1.709921\n",
        "}\n",
        "\n",
        "EPSILON = 1e-50\n",
        "\n",
        "def e_x_lda_unpolarized(rho):\n",
        "  \"\"\"Evaluates LDA exchange energy density for spin unpolarized case.\n",
        "\n",
        "  Parr \u0026 Yang Density-functional theory of atoms and molecules Eq. 6.1.20.\n",
        "\n",
        "  Args:\n",
        "    rho: Float numpy array with shape (num_grids,), the electron density.\n",
        "\n",
        "  Returns:\n",
        "    Float numpy array with shape (num_grids,), the LDA exchange energy density.\n",
        "  \"\"\"\n",
        "  return -3 / 4 * (3 / jnp.pi) ** (1 / 3) * rho ** (4 / 3)\n",
        "\n",
        "\n",
        "def e_x_lda_polarized(rhoa, rhob):\n",
        "  \"\"\"Evaluates LDA exchange energy density for spin polarized case.\n",
        "\n",
        "  Parr \u0026 Yang Density-functional theory of atoms and molecules Eq. 8.2.18.\n",
        "\n",
        "  Args:\n",
        "    rhoa: Float numpy array with shape (num_grids,), the spin up\n",
        "      electron density.\n",
        "    rhob: Float numpy array with shape (num_grids,), the spin down\n",
        "      electron density.\n",
        "\n",
        "  Returns:\n",
        "    Float numpy array with shape (num_grids,), the LDA exchange energy density.\n",
        "  \"\"\"\n",
        "  rho = rhoa + rhob\n",
        "  zeta = (rhoa - rhob) / (rho + EPSILON)  # spin polarization\n",
        "  spin_scaling = 0.5 * ((1 + zeta) ** (4 / 3) + (1 - zeta) ** (4 / 3))\n",
        "  return - 3 / 4 * (3 / jnp.pi) ** (1 / 3) * rho ** (4 / 3) * spin_scaling\n",
        "\n",
        "\n",
        "def g_lda_pw(rs, a, alpha1, beta1, beta2, beta3, beta4):\n",
        "  \"\"\"Evaluates auxiliary function G in the PW parametrization of LDA functional.\n",
        "\n",
        "  10.1103/PhysRevB.45.13244 Eq. 10.\n",
        "\n",
        "  Args:\n",
        "    rs: Float numpy array with shape (num_grids,), Wigner-Seitz radius.\n",
        "    a: Float, parameter.\n",
        "    alpha1: Float, parameter.\n",
        "    beta1: Float, parameter.\n",
        "    beta2: Float, parameter.\n",
        "    beta3: Float, parameter.\n",
        "    beta4: Float, parameter.\n",
        "    use_jax: Boolean, if True, use jax.numpy for calculations, otherwise use\n",
        "      numpy.\n",
        "\n",
        "  Returns:\n",
        "    Float numpy array with shape (num_grids,), auxiliary function G.\n",
        "  \"\"\"\n",
        "  den = 2 * a * (\n",
        "      beta1 * rs**(1 / 2) + beta2 * rs + beta3 * rs**(3 / 2) + beta4 * rs**2)\n",
        "  return -2 * a * (1 + alpha1 * rs) * jnp.log(1 + 1 / den)\n",
        "\n",
        "\n",
        "def get_wigner_seitz_radius(rho):\n",
        "  \"\"\"Evaluates Wigner-Seitz radius of given density.\n",
        "\n",
        "  Args:\n",
        "    rho: Float numpy array with shape (num_grids,), the electron density.\n",
        "\n",
        "  Returns:\n",
        "    Float numpy array with shape (num_grids,), the Wigner-Seitz radius.\n",
        "  \"\"\"\n",
        "  return (3 / (4 * jnp.pi)) ** (1/3) * (rho + EPSILON) ** (-1 / 3)\n",
        "\n",
        "\n",
        "def e_c_lda_unpolarized(rho):\n",
        "  \"\"\"Evaluates LDA correlation energy density for spin unpolarized case.\n",
        "\n",
        "  PW parametrization. 10.1103/PhysRevB.45.13244 Eq. 8-9.\n",
        "\n",
        "  Args:\n",
        "    rho: Float numpy array with shape (num_grids,), the electron density.\n",
        "    use_pbe_params: Boolean, whether use PBE parameters.\n",
        "    use_jax: Boolean, if True, use jax.numpy for calculations, otherwise use\n",
        "      numpy.\n",
        "\n",
        "  Returns:\n",
        "    Float numpy array with shape (num_grids,), the LDA correlation energy\n",
        "      density.\n",
        "  \"\"\"\n",
        "  rs = get_wigner_seitz_radius(rho)\n",
        "  return rho * g_lda_pw(rs=rs, **_LDA_PW_PARAMETERS['eps_c_unpolarized'])\n",
        "\n",
        "\n",
        "def e_c_lda_polarized(rhoa, rhob):\n",
        "  \"\"\"Evaluates LDA correlation energy density for spin polarized case.\n",
        "\n",
        "  PW parametrization. 10.1103/PhysRevB.45.13244 Eq. 8-9.\n",
        "\n",
        "  Args:\n",
        "    rhoa: Float numpy array with shape (num_grids,), the spin up\n",
        "      electron density.\n",
        "    rhob: Float numpy array with shape (num_grids,), the spin down\n",
        "      electron density.\n",
        "    use_pbe_params: Boolean, whether use PBE parameters.\n",
        "\n",
        "  Returns:\n",
        "    Float numpy array with shape (num_grids,), the LDA correlation energy\n",
        "      density.\n",
        "  \"\"\"\n",
        "  rho = rhoa + rhob\n",
        "  rs = get_wigner_seitz_radius(rho)\n",
        "  zeta = (rhoa - rhob) / (rho + EPSILON)\n",
        "\n",
        "  # spin dependent interpolation coefficient\n",
        "  f_zeta = 1 / (2 ** (4 / 3) - 2) * (\n",
        "      (1 + zeta) ** (4 / 3) + (1 - zeta) ** (4 / 3) - 2)\n",
        "\n",
        "  eps_c_unpolarized = g_lda_pw(\n",
        "      rs=rs, **_LDA_PW_PARAMETERS['eps_c_unpolarized'])\n",
        "  eps_c_polarized = g_lda_pw(\n",
        "      rs=rs, **_LDA_PW_PARAMETERS['eps_c_polarized'])\n",
        "  alpha_c = -g_lda_pw(rs=rs, **_LDA_PW_PARAMETERS['alpha_c'])\n",
        "\n",
        "  return rho * (\n",
        "      eps_c_unpolarized\n",
        "      + (1 / _LDA_PW_PARAMETERS['fpp_0']) * alpha_c * f_zeta * (1 - zeta ** 4)\n",
        "      + (eps_c_polarized - eps_c_unpolarized) * f_zeta * zeta ** 4\n",
        "      )\n",
        "\n",
        "\n",
        "def decomposed_e_c_lda_unpolarized(rho):\n",
        "  \"\"\"Evaluates LDA e_c decomposed into same-spin and opposite-spin components.\n",
        "\n",
        "  This function returns the LDA correlation energy density partitioned into\n",
        "  same-spin and opposite-spin components. 10.1063/1.475007 Eq. 7-8.\n",
        "\n",
        "  Args:\n",
        "    rho: Float numpy array with shape (num_grids,), the electron density.\n",
        "\n",
        "  Returns:\n",
        "    e_c_ss: Float numpy array with shape (num_grids,), the same-spin\n",
        "      component of LDA correlation energy density.\n",
        "    e_c_os: Float numpy array with shape (num_grids,), the opposite-spin\n",
        "      component of LDA correlation energy density.\n",
        "  \"\"\"\n",
        "  e_c = e_c_lda_unpolarized(rho)\n",
        "  e_c_ss = 2 * e_c_lda_polarized(rho / 2, jnp.zeros_like(rho))\n",
        "  e_c_os = e_c - e_c_ss\n",
        "\n",
        "  return e_c_ss, e_c_os\n",
        "\n",
        "\n",
        "def decomposed_e_c_lda_polarized(rhoa, rhob):\n",
        "  \"\"\"Evaluates LDA e_c decomposed into same-spin and opposite-spin components.\n",
        "\n",
        "  This function returns the LDA correlation energy density partitioned into\n",
        "  same-spin and opposite-spin components. 10.1063/1.475007 Eq. 7-8.\n",
        "\n",
        "  Args:\n",
        "    rhoa: Float numpy array with shape (num_grids,), the spin up\n",
        "      electron density.\n",
        "    rhob: Float numpy array with shape (num_grids,), the spin down\n",
        "      electron density.\n",
        "\n",
        "  Returns:\n",
        "    e_c_aa: Float numpy array with shape (num_grids,), the same-spin (aa)\n",
        "      component of LDA correlation energy density.\n",
        "    e_c_bb: Float numpy array with shape (num_grids,), the same-spin (bb)\n",
        "      component of LDA correlation energy density.\n",
        "    e_c_ab: Float numpy array with shape (num_grids,), the opposite-spin\n",
        "      component of LDA correlation energy density.\n",
        "  \"\"\"\n",
        "  zero = jnp.zeros_like(rhoa)\n",
        "\n",
        "  e_c = e_c_lda_polarized(rhoa, rhob)\n",
        "  e_c_aa = e_c_lda_polarized(rhoa, zero)\n",
        "  e_c_bb = e_c_lda_polarized(zero, rhob)\n",
        "  e_c_ab = e_c - e_c_aa - e_c_bb\n",
        "\n",
        "  return e_c_aa, e_c_bb, e_c_ab\n",
        "\n",
        "\n",
        "def get_reduced_density_gradient(rho, sigma):\n",
        "  \"\"\"Evaluates reduced density gradient.\n",
        "\n",
        "  Args:\n",
        "    rho: Float numpy array with shape (num_grids,), the electron density.\n",
        "    sigma: Float numpy array with shape (num_grids,), the norm square of\n",
        "      density gradient.\n",
        "\n",
        "  Returns:\n",
        "    Float numpy array with shape (num_grids,), the reduced density gradient.\n",
        "  \"\"\"\n",
        "  # NOTE: EPSILON is inserted in such a way that the reduce density\n",
        "  # gradient is zero if rho or sigma or both are zero.\n",
        "  return jnp.sqrt(sigma + EPSILON ** 3) / (rho ** (4 / 3) + EPSILON)\n",
        "\n",
        "\n",
        "def get_mgga_t(rho, tau, polarized):\n",
        "  \"\"\"Evaluates the auxiliary quantity t in meta-GGA functional forms.\n",
        "\n",
        "  t = (tau_HEG / tau),  where tau_HEG is the kinetic energy density of\n",
        "  homogeneous electron gass.\n",
        "\n",
        "  Args:\n",
        "    rho: Float numpy array with shape (num_grids,), the electron density.\n",
        "    tau: Float numpy array with shape (num_grids,), the kinetic energy density.\n",
        "    polarized: Boolean, whether the system is spin polarized.\n",
        "\n",
        "  Returns:\n",
        "    Float numpy array with shape (num_grids,), the auxiliary quantity w.\n",
        "  \"\"\"\n",
        "  spin_factor = 1 if polarized else (1 / 2) ** (2 / 3)\n",
        "  # 3 / 10 instead of 3 / 5 is used below because tau is defined with 1 / 2\n",
        "  tau_heg = 3 / 10 * (6 * jnp.pi ** 2) ** (2 / 3) * rho ** (5 / 3)\n",
        "  return spin_factor * tau_heg / (tau + EPSILON)\n",
        "\n",
        "\n",
        "def f_rsh(rho, omega=RSH_PARAMS[0], polarized=False,):\n",
        "  \"\"\"Enchancement factor for evaluating short-range semilocal exchange.\n",
        "\n",
        "  10.1063/1.4952647 Eq. 11.\n",
        "\n",
        "  Args:\n",
        "    rho: Float numpy array with shape (num_grids,), the electron density.\n",
        "    omega: Float, the range seperation parameter.\n",
        "    polarized: Boolean, whether the system is spin polarized.\n",
        "\n",
        "  Returns:\n",
        "    Float numpy array with shape (num_grids,), the RSH enhancement factor.\n",
        "  \"\"\"\n",
        "  spin_factor = 1 if polarized else 2\n",
        "  # Fermi wave vector\n",
        "  kf = (6 * jnp.pi**2 * rho / spin_factor + EPSILON) ** (1 / 3)\n",
        "  a = omega / kf + EPSILON  # variable a in Eq. 11\n",
        "  return (1 - 2 / 3 * a * (2 * jnp.pi ** (1 / 2) * jax.scipy.special.erf(1 / a) - 3 * a\n",
        "                           + a ** 3 + (2 * a - a ** 3) * jnp.exp(-1 / a ** 2)))\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "MTkaZzKN4ZlD"
      },
      "source": [
        "#GAS22 enhancement factors\n",
        "The following enhancement factors for the GAS22 functional corresponds to the Eq. 9-11 in the manuscript."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "hSnOjb0z4c16"
      },
      "outputs": [],
      "source": [
        "# NOTE: In the following functions, the square of reduced density gradient are\n",
        "# used as function arguments (x2).\n",
        "\n",
        "def get_u(x2, gamma):\n",
        "  \"\"\"Evaluates the auxiliary quantity u = gamma x^2 / (1 + gamma x^2).\"\"\"\n",
        "  return gamma * x2 / (1 + gamma * x2)\n",
        "\n",
        "def f_x_gas22(x2, w):\n",
        "  \"\"\"Evaluates the exchange enhancement factor for the GAS22 functional.\"\"\"\n",
        "  u = get_u(x2, gamma=0.003840616724010807)\n",
        "  return 0.862139736374172 + 0.936993691972698 * u  + 0.317533683085033 * w\n",
        "\n",
        "def f_css_gas22(x2, w):\n",
        "  \"\"\"Evaluates the same-spin correlation enhancement factor for the GAS22 functional.\"\"\"\n",
        "  u = get_u(x2, gamma=0.46914023462026644)\n",
        "  return (u - 4.10753796482853 * w - 5.24218990333846 * w**2\n",
        "          - 1.76643208454076 * u**6 + 7.5380689617542 * u**6 * w**4)\n",
        "\n",
        "def f_cos_gas22(x2, w):\n",
        "  \"\"\"Evaluates the opposite-spin correlation enhancement factor for the GAS22 functional.\"\"\"\n",
        "  return (0.805124374375355 + 7.98909430970845 * w**2\n",
        "          - 1.76098915061634 * w**2 * jnp.cbrt(x2) - 7.54815900595292 * w**6\n",
        "          + 2.00093961824784 * w**6 * jnp.cbrt(x2))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Gh8yLDQMraoe"
      },
      "source": [
        "# Define GAS22 in the format of PySCF custom functions"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "WB4vfZ9RljNH"
      },
      "outputs": [],
      "source": [
        "def gas22_unpolarized(rho, sigma, tau):\n",
        "  \"\"\"Evaluates XC energy density for spin unpolarized case.\n",
        "\n",
        "  Args:\n",
        "    rho: Float numpy array with shape (num_grids,), the electron density.\n",
        "    sigma: Float numpy array with shape (num_grids,), the norm square of\n",
        "      density gradient.\n",
        "    tau: Float numpy array with shape (num_grids,), the kinetic energy\n",
        "      density.\n",
        "\n",
        "  Returns:\n",
        "    Float numpy array with shape (num_grids,), the XC energy density.\n",
        "  \"\"\"\n",
        "  rho_s = 0.5 * rho\n",
        "  x_s = get_reduced_density_gradient(rho_s, 0.25 * sigma)\n",
        "  x2_s = x_s ** 2\n",
        "  t_s = get_mgga_t(rho_s, 0.5 * tau, polarized=True)\n",
        "  w_s = (t_s - 1) / (t_s + 1)\n",
        "\n",
        "  e_lda_x = e_x_lda_unpolarized(rho) * f_rsh(rho, polarized=False)\n",
        "  e_lda_css, e_lda_cos = decomposed_e_c_lda_unpolarized(rho)\n",
        "  \n",
        "  f_x = f_x_gas22(x2_s, w_s)\n",
        "  f_css = f_css_gas22(x2_s, w_s)\n",
        "  f_cos = f_cos_gas22(x2_s, w_s)\n",
        "\n",
        "  return e_lda_x * f_x + e_lda_css * f_css + e_lda_cos * f_cos\n",
        "\n",
        "def gas22_polarized(rho_a, rho_b, sigma_aa, sigma_ab, sigma_bb, tau_a, tau_b):\n",
        "  \"\"\"Evaluates XC energy density for spin polarized case.\n",
        "\n",
        "  Args:\n",
        "    rho_a: Float numpy array with shape (num_grids,), the spin up electron\n",
        "      density.\n",
        "    rho_b: Float numpy array with shape (num_grids,), the spin down electron\n",
        "      density.\n",
        "    sigma_aa: Float numpy array with shape (num_grids,), the norm square of\n",
        "      density gradient (aa component).\n",
        "    sigma_ab: Float numpy array with shape (num_grids,), the norm square of\n",
        "      density gradient (ab component).\n",
        "    sigma_bb: Float numpy array with shape (num_grids,), the norm square of\n",
        "      density gradient (bb component).\n",
        "    tau_a: Float numpy array with shape (num_grids,), the spin up kinetic\n",
        "      energy density.\n",
        "    tau_b: Float numpy array with shape (num_grids,), the spin down kinetic\n",
        "      energy density.\n",
        "\n",
        "  Returns:\n",
        "    Float numpy array with shape (num_grids,), the XC energy density.\n",
        "  \"\"\"\n",
        "  del sigma_ab\n",
        "  rho_ab = 0.5 * (rho_a + rho_b)\n",
        "\n",
        "  x_a = get_reduced_density_gradient(rho_a, sigma_aa)\n",
        "  x_b = get_reduced_density_gradient(rho_b, sigma_bb)\n",
        "  x2_a = x_a ** 2\n",
        "  x2_b = x_b ** 2\n",
        "  x2_ab = 0.5 * (x2_a + x2_b)\n",
        "\n",
        "  t_a = get_mgga_t(rho_a, tau_a, polarized=True)\n",
        "  t_b = get_mgga_t(rho_b, tau_b, polarized=True)\n",
        "  t_ab = 0.5 * (t_a + t_b)\n",
        "  w_a = (t_a - 1) / (t_a + 1)\n",
        "  w_b = (t_b - 1) / (t_b + 1)\n",
        "  w_ab = (t_ab - 1) / (t_ab + 1)\n",
        "\n",
        "  e_lda_x_a = 0.5 * e_x_lda_unpolarized(2 * rho_a) * f_rsh(\n",
        "      rho_a, polarized=True)\n",
        "  e_lda_x_b = 0.5 * e_x_lda_unpolarized(2 * rho_b) * f_rsh(\n",
        "      rho_b, polarized=True)\n",
        "  e_lda_css_a, e_lda_css_b, e_lda_cos = decomposed_e_c_lda_polarized(\n",
        "      rho_a, rho_b)\n",
        "\n",
        "  f_x_a = f_x_gas22(x2_a, w_a)\n",
        "  f_x_b = f_x_gas22(x2_b, w_b)\n",
        "  f_css_a = f_css_gas22(x2_a, w_a)\n",
        "  f_css_b = f_css_gas22(x2_b, w_b)\n",
        "  f_cos = f_cos_gas22(x2_ab, w_ab)\n",
        "\n",
        "  return (e_lda_x_a * f_x_a + e_lda_x_b * f_x_b\n",
        "          + e_lda_css_a * f_css_a + e_lda_css_b * f_css_b + e_lda_cos * f_cos)\n",
        "\n",
        "def eval_xc_gas22(xc_code,\n",
        "                  rho_and_derivs,\n",
        "                  spin=0,\n",
        "                  relativity=0,\n",
        "                  deriv=1,\n",
        "                  verbose=None):\n",
        "  \"\"\"Evaluates exchange-correlation energy densities and derivatives.\n",
        "\n",
        "  Args:\n",
        "    xc_code: A dummy argument, not used.\n",
        "    rho_and_derivs: Float numpy array with shape (6, num_grids) for spin\n",
        "      unpolarized case; 2-tuple of float numpy array with shape (6, num_grids)\n",
        "      for spin polarized case. Electron density and its derivatives. For\n",
        "      spin unpolarized case, the 6 subarrays represent (density, gradient_x,\n",
        "      gradient_y, gradient_z, laplacian, tau); for spin polarized case, the\n",
        "      spin up and spin down densities and derivatives are each represented\n",
        "      with a (6, num_grids) array.\n",
        "    spin: Integer, 0 for spin unpolarized and 1 for spin polarized\n",
        "      calculations.\n",
        "    relativity: A dummy argument, not used.\n",
        "    deriv: Integer, the order of derivatives evaluated for XC energy density.\n",
        "    verbose: A dummy argument, not used.\n",
        "\n",
        "  Returns:\n",
        "    eps_xc: Float numpy array with shape (num_grids,), the XC energy density\n",
        "      per particle.\n",
        "    v_xc: Tuple of float numpy arrays, the first derivatives of XC energy\n",
        "      density per volume to various quantities. See pyscf/dft/libxc.py for\n",
        "      more details.\n",
        "    f_xc: A dummy return value, not used.\n",
        "    k_xc: A dummy return value, not used.\n",
        "\n",
        "  Raises:\n",
        "    NotImplementedError: If derivative order higher than one is requested\n",
        "      (deriv \u003e 1).\n",
        "  \"\"\"\n",
        "  del xc_code, relativity, verbose\n",
        "\n",
        "  if deriv != 1:\n",
        "    raise NotImplementedError('Only deriv = 1 is implemented.')\n",
        "\n",
        "  if spin == 0:\n",
        "    rho, grad_x, grad_y, grad_z, _, tau = rho_and_derivs\n",
        "    sigma = grad_x**2 + grad_y**2 + grad_z**2\n",
        "\n",
        "    e_xc, grads = jax.jit(jax.vmap(jax.value_and_grad(\n",
        "        gas22_unpolarized, argnums=(0, 1, 2))))(rho, sigma, tau)\n",
        "    vrho, vsigma, vtau = grads\n",
        "\n",
        "  else:\n",
        "    rhoa, grad_x_a, grad_y_a, grad_z_a, _, tau_a = rho_and_derivs[0]\n",
        "    rhob, grad_x_b, grad_y_b, grad_z_b, _, tau_b = rho_and_derivs[1]\n",
        "    rho = rhoa + rhob\n",
        "    sigma_aa = grad_x_a**2 + grad_y_a**2 + grad_z_a**2\n",
        "    sigma_ab = grad_x_a * grad_x_b + grad_y_a * grad_y_b + grad_z_a * grad_z_b\n",
        "    sigma_bb = grad_x_b**2 + grad_y_b**2 + grad_z_b**2\n",
        "\n",
        "    e_xc, grads = jax.jit(jax.vmap(jax.value_and_grad(\n",
        "        gas22_polarized, argnums=(0, 1, 2, 3, 4, 5, 6))))(\n",
        "            rhoa, rhob, sigma_aa, sigma_ab, sigma_bb, tau_a, tau_b)\n",
        "\n",
        "    vrhoa, vrhob, vsigma_aa, vsigma_ab, vsigma_bb, vtau_a, vtau_b = grads\n",
        "    vrho = np.stack((vrhoa, vrhob), axis=1)\n",
        "    vsigma = np.stack((vsigma_aa, vsigma_ab, vsigma_bb), axis=1)\n",
        "    vtau = np.stack((vtau_a, vtau_b), axis=1)\n",
        "\n",
        "  eps_xc = e_xc / (rho + EPSILON)\n",
        "  return (np.array(eps_xc),\n",
        "          (np.array(vrho), np.array(vsigma),\n",
        "            np.zeros_like(vtau), np.array(vtau)),\n",
        "          None,\n",
        "          None)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "5nVLL7TsCJGX"
      },
      "source": [
        "# SCF calculations with PySCF"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "510maFzNFYBH"
      },
      "outputs": [],
      "source": [
        "def run_scf_for_mol(atom,\n",
        "                    basis,\n",
        "                    charge,\n",
        "                    spin,\n",
        "                    conv_tol=1e-6,\n",
        "                    use_sg1_prune_for_nlc=True):\n",
        "  \"\"\"Runs SCF calculations for molecule using GAS22 functional.\n",
        "  \n",
        "  Args:\n",
        "    atom: String, the atomic structure.\n",
        "    charge: Integer, the total charge of molecule.\n",
        "    spin: Integer, the difference of spin up and spin down electron numbers.\n",
        "    basis: String, the GTO basis.\n",
        "    conv_tol: Float, the convergence threshold of total energy.\n",
        "    use_sg1_prune_for_nlc: Boolean, whether use SG1 prune for NLC calculation.\n",
        "  \n",
        "  Returns:\n",
        "    Float, the SCF total energy.\n",
        "  \"\"\"\n",
        "  mol = gto.M(atom=atom, basis=basis, charge=charge, spin=spin, verbose=4)\n",
        "\n",
        "  if spin == 0:\n",
        "    ks = dft.RKS(mol)\n",
        "  else:\n",
        "    ks = dft.UKS(mol)\n",
        "\n",
        "  ks.define_xc_(eval_xc_gas22, xctype='MGGA', hyb=HYBRID_COEFF, rsh=RSH_PARAMS)\n",
        "\n",
        "  # NOTE: This is necessary because PySCF will use this name to decide the\n",
        "  # parameters for VV10 NLC. GAS22 has identical nonlocal XC as wB97M-V, including\n",
        "  # VV10 NLC. \n",
        "  ks.xc = 'wb97mv'\n",
        "  ks.nlc = 'VV10'\n",
        "\n",
        "  if use_sg1_prune_for_nlc:\n",
        "    # NOTE: SG1 prune can be used to reduce the computational cost of\n",
        "    # VV10 NLC. The use of SG1 prune has very little effect on the resulting\n",
        "    # XC energy. SG1 prune is used in PySCF's example\n",
        "    # pyscf/examples/dft/33-nlc_functionals.py and is also used in paper\n",
        "    # 10.1080/00268976.2017.1333644. Note that SG1 prune in PySCF is not\n",
        "    # available for some elements appeared in the MCGDB84 database.\n",
        "    ks.nlcgrids.prune = dft.gen_grid.sg1_prune\n",
        "\n",
        "  # NOTE: It is necessary to override ks._numint._xc_type method to\n",
        "  # let PySCF correctly use a custom XC functional with NLC. Also, note that\n",
        "  # ks.xc is used to determine NLC parameters.\n",
        "  ks._numint._xc_type = lambda code: 'NLC' if 'VV10' in code else 'MGGA'\n",
        "\n",
        "  ks.conv_tol = conv_tol\n",
        "\n",
        "  ks.kernel()\n",
        "\n",
        "  return ks.e_tot"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "executionInfo": {
          "elapsed": 65921,
          "status": "ok",
          "timestamp": 1644345845579,
          "user": {
            "displayName": "Li Li",
            "photoUrl": "https://lh3.googleusercontent.com/a-/AOh14GjUXuHpz1GDBTLkFU6msl96ZWWDxtMqu7UXHH6qLg=s64",
            "userId": "05582677413017445658"
          },
          "user_tz": 480
        },
        "id": "vXFfaeYyFjae",
        "outputId": "84e9a969-a9aa-4994-a1d8-4f5f3869adf7"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "System: uname_result(system='Linux', node='adbbb022342a', release='5.4.144+', version='#1 SMP Tue Dec 7 09:58:10 PST 2021', machine='x86_64', processor='x86_64')  Threads 2\n",
            "Python 3.7.12 (default, Jan 15 2022, 18:48:18) \n",
            "[GCC 7.5.0]\n",
            "numpy 1.19.5  scipy 1.4.1\n",
            "Date: Tue Feb  8 18:42:59 2022\n",
            "PySCF version 1.7.6\n",
            "PySCF path  /usr/local/lib/python3.7/dist-packages/pyscf\n",
            "\n",
            "[CONFIG] conf_file None\n",
            "[INPUT] verbose = 4\n",
            "[INPUT] num. atoms = 3\n",
            "[INPUT] num. electrons = 10\n",
            "[INPUT] charge = 0\n",
            "[INPUT] spin (= nelec alpha-beta = 2S) = 0\n",
            "[INPUT] symmetry False subgroup None\n",
            "[INPUT] Mole.unit = angstrom\n",
            "[INPUT]  1 O     -0.084889000000   0.056804000000   0.055205000000 AA   -0.160416960988   0.107344002780   0.104322330707 Bohr\n",
            "[INPUT]  2 H      0.710375000000   0.517729000000   0.413971000000 AA    1.342414195738   0.978366016745   0.782291813512 Bohr\n",
            "[INPUT]  3 H     -0.890958000000   0.498588000000   0.414274000000 AA   -1.683666608490   0.942194768995   0.782864400528 Bohr\n",
            "\n",
            "nuclear repulsion = 8.91072304254368\n",
            "number of shells = 42\n",
            "number of NR pGTOs = 158\n",
            "number of NR cGTOs = 132\n",
            "basis = def2qzvppd\n",
            "ecp = {}\n",
            "CPU time:         4.02\n",
            "\n",
            "\n",
            "******** \u003cclass 'pyscf.dft.rks.RKS'\u003e ********\n",
            "method = RKS-RHF\n",
            "initial guess = minao\n",
            "damping factor = 0\n",
            "level_shift factor = 0\n",
            "DIIS = \u003cclass 'pyscf.scf.diis.CDIIS'\u003e\n",
            "diis_start_cycle = 1\n",
            "diis_space = 8\n",
            "SCF conv_tol = 1e-06\n",
            "SCF conv_tol_grad = None\n",
            "SCF max_cycles = 50\n",
            "direct_scf = True\n",
            "direct_scf_tol = 1e-13\n",
            "chkfile to save SCF result = /content/tmp3tccv4ww\n",
            "max_memory 4000 MB (current use 266 MB)\n",
            "XC functionals = wb97mv\n",
            "NLC functional = VV10\n",
            "small_rho_cutoff = 1e-07\n",
            "radial grids: \n",
            "    Treutler-Ahlrichs [JCP 102, 346 (1995); DOI:10.1063/1.469408] (M4) radial grids\n",
            "    \n",
            "becke partition: Becke, JCP 88, 2547 (1988); DOI:10.1063/1.454033\n",
            "pruning grids: \u003cfunction nwchem_prune at 0x7fd3616c69e0\u003e\n",
            "grids dens level: 3\n",
            "symmetrized grids: False\n",
            "atomic radii adjust function: \u003cfunction treutler_atomic_radii_adjust at 0x7fd3616c6e60\u003e\n",
            "** Following is NLC Grids **\n",
            "radial grids: \n",
            "    Treutler-Ahlrichs [JCP 102, 346 (1995); DOI:10.1063/1.469408] (M4) radial grids\n",
            "    \n",
            "becke partition: Becke, JCP 88, 2547 (1988); DOI:10.1063/1.454033\n",
            "pruning grids: \u003cfunction sg1_prune at 0x7fd3616bc7a0\u003e\n",
            "grids dens level: 3\n",
            "symmetrized grids: False\n",
            "atomic radii adjust function: \u003cfunction treutler_atomic_radii_adjust at 0x7fd3616c6e60\u003e\n",
            "Set gradient conv threshold to 0.001\n",
            "tot grids = 34310\n",
            "tot grids = 15258\n"
          ]
        },
        {
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "WARNING:absl:No GPU/TPU found, falling back to CPU. (Set TF_CPP_MIN_LOG_LEVEL=0 and rerun for more info.)\n"
          ]
        },
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "init E= -76.2460632565618\n",
            "  HOMO = -0.45183892088776  LUMO = -0.000563452226807279\n",
            "cycle= 1 E= -76.3638063214408  delta_E= -0.118  |g|= 0.593  |ddm|= 0.726\n",
            "  HOMO = -0.291693245914626  LUMO = 0.0261752722564711\n",
            "cycle= 2 E= -76.3554305875159  delta_E= 0.00838  |g|= 0.678  |ddm|= 0.472\n",
            "  HOMO = -0.42835098458092  LUMO = 0.0229899941959564\n",
            "cycle= 3 E= -76.4299324142884  delta_E= -0.0745  |g|= 0.0278  |ddm|= 0.207\n",
            "  HOMO = -0.424557873317824  LUMO = 0.025502365859663\n",
            "cycle= 4 E= -76.4300487284296  delta_E= -0.000116  |g|= 0.00708  |ddm|= 0.0187\n",
            "  HOMO = -0.425933919786901  LUMO = 0.0254451196686611\n",
            "cycle= 5 E= -76.4300600381892  delta_E= -1.13e-05  |g|= 0.000857  |ddm|= 0.00886\n",
            "  HOMO = -0.426038301180972  LUMO = 0.0254047853829936\n",
            "cycle= 6 E= -76.4300602793656  delta_E= -2.41e-07  |g|= 9.33e-05  |ddm|= 0.00212\n",
            "  HOMO = -0.426060779696682  LUMO = 0.0254012851301158\n",
            "Extra cycle  E= -76.4300602830018  delta_E= -3.64e-09  |g|= 5.55e-05  |ddm|= 0.000322\n",
            "converged SCF energy = -76.4300602830018\n"
          ]
        },
        {
          "data": {
            "text/plain": [
              "-76.43006028300178"
            ]
          },
          "execution_count": 7,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "# Example: water molecule (spin unpolarized)\n",
        "# Atomic structure is taken from 01a_water_monA_3B-69.xyz in MGCDB84\n",
        "# Expected Etot: 76.430060 au\n",
        "run_scf_for_mol(\n",
        "    atom = \"\"\"\\\n",
        "    O -0.0848890000 0.0568040000 0.0552050000\n",
        "    H 0.7103750000 0.5177290000 0.4139710000\n",
        "    H -0.8909580000 0.4985880000 0.4142740000\n",
        "    \"\"\",\n",
        "    basis = 'def2qzvppd',\n",
        "    charge = 0,\n",
        "    spin = 0,\n",
        ")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "executionInfo": {
          "elapsed": 64420,
          "status": "ok",
          "timestamp": 1644345909978,
          "user": {
            "displayName": "Li Li",
            "photoUrl": "https://lh3.googleusercontent.com/a-/AOh14GjUXuHpz1GDBTLkFU6msl96ZWWDxtMqu7UXHH6qLg=s64",
            "userId": "05582677413017445658"
          },
          "user_tz": 480
        },
        "id": "G3eKMuMhF4qW",
        "outputId": "1d178e6c-3f34-432f-8ddd-a26af76780a8"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "System: uname_result(system='Linux', node='adbbb022342a', release='5.4.144+', version='#1 SMP Tue Dec 7 09:58:10 PST 2021', machine='x86_64', processor='x86_64')  Threads 2\n",
            "Python 3.7.12 (default, Jan 15 2022, 18:48:18) \n",
            "[GCC 7.5.0]\n",
            "numpy 1.19.5  scipy 1.4.1\n",
            "Date: Tue Feb  8 18:44:05 2022\n",
            "PySCF version 1.7.6\n",
            "PySCF path  /usr/local/lib/python3.7/dist-packages/pyscf\n",
            "\n",
            "[CONFIG] conf_file None\n",
            "[INPUT] verbose = 4\n",
            "[INPUT] num. atoms = 1\n",
            "[INPUT] num. electrons = 15\n",
            "[INPUT] charge = 0\n",
            "[INPUT] spin (= nelec alpha-beta = 2S) = 3\n",
            "[INPUT] symmetry False subgroup None\n",
            "[INPUT] Mole.unit = angstrom\n",
            "[INPUT]  1 P      0.000000000000   0.000000000000   0.000000000000 AA    0.000000000000   0.000000000000   0.000000000000 Bohr\n",
            "\n",
            "nuclear repulsion = 0\n",
            "number of shells = 24\n",
            "number of NR pGTOs = 111\n",
            "number of NR cGTOs = 76\n",
            "basis = def2qzvppd\n",
            "ecp = {}\n",
            "CPU time:       103.62\n",
            "\n",
            "\n",
            "******** \u003cclass 'pyscf.dft.uks.UKS'\u003e ********\n",
            "method = UKS-UHF\n",
            "initial guess = minao\n",
            "damping factor = 0\n",
            "level_shift factor = 0\n",
            "DIIS = \u003cclass 'pyscf.scf.diis.CDIIS'\u003e\n",
            "diis_start_cycle = 1\n",
            "diis_space = 8\n",
            "SCF conv_tol = 1e-06\n",
            "SCF conv_tol_grad = None\n",
            "SCF max_cycles = 50\n",
            "direct_scf = True\n",
            "direct_scf_tol = 1e-13\n",
            "chkfile to save SCF result = /content/tmp2hqsgpp6\n",
            "max_memory 4000 MB (current use 1118 MB)\n",
            "number electrons alpha = 9  beta = 6\n",
            "XC functionals = wb97mv\n",
            "NLC functional = VV10\n",
            "small_rho_cutoff = 1e-07\n",
            "radial grids: \n",
            "    Treutler-Ahlrichs [JCP 102, 346 (1995); DOI:10.1063/1.469408] (M4) radial grids\n",
            "    \n",
            "becke partition: Becke, JCP 88, 2547 (1988); DOI:10.1063/1.454033\n",
            "pruning grids: \u003cfunction nwchem_prune at 0x7fd3616c69e0\u003e\n",
            "grids dens level: 3\n",
            "symmetrized grids: False\n",
            "atomic radii adjust function: \u003cfunction treutler_atomic_radii_adjust at 0x7fd3616c6e60\u003e\n",
            "** Following is NLC Grids **\n",
            "radial grids: \n",
            "    Treutler-Ahlrichs [JCP 102, 346 (1995); DOI:10.1063/1.469408] (M4) radial grids\n",
            "    \n",
            "becke partition: Becke, JCP 88, 2547 (1988); DOI:10.1063/1.454033\n",
            "pruning grids: \u003cfunction sg1_prune at 0x7fd3616bc7a0\u003e\n",
            "grids dens level: 3\n",
            "symmetrized grids: False\n",
            "atomic radii adjust function: \u003cfunction treutler_atomic_radii_adjust at 0x7fd3616c6e60\u003e\n",
            "Set gradient conv threshold to 0.001\n",
            "tot grids = 18880\n",
            "tot grids = 6420\n",
            "init E= -340.984264138224\n",
            "  alpha nocc = 9  HOMO = -0.184697973539261  LUMO = 0.042553321200819\n",
            "  beta  nocc = 6  HOMO = -0.640396329357465  LUMO = -0.184697973539265\n",
            "\n",
            "WARN: system HOMO -0.184697973539261 \u003e= system LUMO -0.184697973539265\n",
            "\n",
            "cycle= 1 E= -341.243196710095  delta_E= -0.259  |g|= 0.14  |ddm|= 0.632\n",
            "  alpha nocc = 9  HOMO = -0.371473682197285  LUMO = 0.0509025039179709\n",
            "  beta  nocc = 6  HOMO = -0.624190245086803  LUMO = -0.0990119234781705\n",
            "cycle= 2 E= -341.253001324628  delta_E= -0.0098  |g|= 0.0604  |ddm|= 0.282\n",
            "  alpha nocc = 9  HOMO = -0.350798261377467  LUMO = 0.0526944485815933\n",
            "  beta  nocc = 6  HOMO = -0.590501167341494  LUMO = -0.0739532149711042\n",
            "cycle= 3 E= -341.253807848467  delta_E= -0.000807  |g|= 0.00977  |ddm|= 0.0691\n",
            "  alpha nocc = 9  HOMO = -0.346091008379461  LUMO = 0.0527073811670932\n",
            "  beta  nocc = 6  HOMO = -0.579896938821247  LUMO = -0.0616632725597869\n",
            "cycle= 4 E= -341.253897789897  delta_E= -8.99e-05  |g|= 0.00158  |ddm|= 0.029\n",
            "  alpha nocc = 9  HOMO = -0.346059822769328  LUMO = 0.0524962553174933\n",
            "  beta  nocc = 6  HOMO = -0.579242442999468  LUMO = -0.0604496702532998\n",
            "cycle= 5 E= -341.253900641421  delta_E= -2.85e-06  |g|= 0.000232  |ddm|= 0.00527\n",
            "  alpha nocc = 9  HOMO = -0.34598577148016  LUMO = 0.0525046230228023\n",
            "  beta  nocc = 6  HOMO = -0.5791784512025  LUMO = -0.0605782504060349\n",
            "cycle= 6 E= -341.253900677  delta_E= -3.56e-08  |g|= 2.1e-05  |ddm|= 0.000885\n",
            "  alpha nocc = 9  HOMO = -0.346002996809474  LUMO = 0.0525127001322863\n",
            "  beta  nocc = 6  HOMO = -0.579198566496345  LUMO = -0.0606008850014173\n",
            "Extra cycle  E= -341.253900677469  delta_E= -4.69e-10  |g|= 4.6e-06  |ddm|= 7.09e-05\n",
            "converged SCF energy = -341.253900677469  \u003cS^2\u003e = 3.752221  2S+1 = 4.0011103\n"
          ]
        },
        {
          "data": {
            "text/plain": [
              "-341.25390067746855"
            ]
          },
          "execution_count": 8,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "# Example: P atom (spin polarized)\n",
        "# Atomic structure is taken from 25_P_AE18.xyz in MGCDB84\n",
        "# Expected Etot: -341.253901 au\n",
        "run_scf_for_mol(\n",
        "    atom = \"P 0.0000000000 0.0000000000 0.0000000000\",\n",
        "    basis = 'def2qzvppd',\n",
        "    charge = 0,\n",
        "    spin = 3,\n",
        ")"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "collapsed_sections": [
        "Bkq_0Rw8G9WY",
        "hEFvWs_YB1qE",
        "Gh8yLDQMraoe"
      ],
      "name": "run_GAS22.ipynb",
      "provenance": [
        {
          "file_id": "1IQeJJ9SkicUwwHR15moUEg9cfX67yUIj",
          "timestamp": 1644347977477
        },
        {
          "file_id": "1fLegX-tukvOz0GKwBGGlu5VBXjITq4Me",
          "timestamp": 1644345717953
        }
      ]
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}
